# youyi
# -*- coding: utf-8 -*-

"""
FE(Feedback Evolution) 算法实现

流程：
    1. 生成初始提示词种群
    2. 评估初始提示词种群
    3. 使用遗传算法进行进化
        3.1 选择亲代
        3.2 抽取样本
        3.3 生成反馈
        3.4 生成子代
        3.5 评估子代
        3.6 更新种群
    4. 评估找到最佳提示词

会把每次的结果放入同目录下的“simplified_FE_results”文件夹中
    以运行时间命名该文件夹
    info.txt: 本次运行的基本情况
    results.json: 本次运行的中间结果
    results_val.json: 本次运行的结果（测试集）
    eval_config.json: 评估参数配置
    eval_optimized_config.json: 评估优化配置

    
集成了独立的评估框架 youyi_evaluator
评估及优化参数说明:
    eval_config基本参数
        evaluation_index: 评估指标(acc, bleu, rouge, BERTscore, LLM_mark)
        carry_template: 执行模板（用于让执行模型生成），包含[PROMPT], [INPUT], [OUTPUT]标记
        carry_matcher: 执行模式匹配器，用于解析执行模型的返回
        carry_model: 执行模型
    eval_config特定参数
        language: 语言，不同语言的分词方法不同(en, cn)，关联指标: bleu, rouge, BERTscore
        ROUGE_kind: 类型, str, 计算方式, 主要变体有ROUGE-n与ROUGE-L ["1", "2", "3", "4", "5", "l"]，关联指标: rouge
        ROUGE_kind2: 类型2, str, 具体指标 ["r", "p", "f"]，关联指标: rouge
        BERTscore_kind: 类型, str, 具体指标 ["P", "R", "F1"]，关联指标: BERTscore

        evaluation_template: 评估模板，关联指标: LLM_mark
        evaluation_matcher: 评估模式匹配器，用于解析评估模型的返回，关联指标: LLM_mark
        evaluation_models: 评估模型（用于打分式评估，代理专家）, list, 元素为评估模型实例，关联指标: LLM_mark
        task_kind: 问题类型, str, ['封闭', '开放']，关联指标: LLM_mark
        consistency_judgment: str, ["cohen_kappa", "cohen_kappa_linear", "cohen_kappa_quadratic", "ICC", None], 一致性判断方法，关联指标: LLM_mark
            评估标准较多时，建议采用 cohen_kappa_quadratic, ICC; 评估标准较少时，建议采用 cohen_kappa, cohen_kappa_linear
        consistency_threshold: float, 一致性判断方法的阈值，推荐[0.2, 0.2, 0.2, 0.4, None]，关联指标: LLM_mark
"""

import random
import re
from typing import List, Tuple, Dict, Any, Optional, Callable, Union
from tqdm import tqdm # 命令行进度条
import sys, os, json, datetime
from pprint import pprint
import traceback


# 导入根目录 (my_work 文件夹中)
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# 导入评估（优化）函数、模式匹配器基类
from youyi_evaluator import batch_UCB, BaseMatcher
# 导入大语言模型
from BM import Da_Mo_Xing_2

# 导入根目录 (my_work 文件夹中)
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
# 导入提示词模板
from prompt_templates import PROMPT_generate_TEMPLATE, PROMPT_feedback_TEMPLATE, PROMPT_direct_TEMPLATE, PROMPT_child_TEMPLATE,\
PROMPT_error_patterns_TEMPLATE, PROMPT_root_cause_TEMPLATE, PROMPT_improvement_TEMPLATE, PROMPT_feedback_synthesis_TEMPLATE,\
PROMPT_knowledge_extraction_TEMPLATE



class SimpleFE:
    """简化版FE算法，聚焦核心功能"""
    
    def __init__(self, carry_model, optimize_model, val_model):
        """
        初始化FE算法
        
        参数:
            carry_model: Da_Mo_Xing_2实例，执行提示词的模型，提示词所针对的对象
            optimize_model: Da_Mo_Xing_2实例，优化提示词的模型
            val_model: Da_Mo_Xing_2实例，用于评估提示词的模型，通过比较二者输出
            
        返回:
            None
        """
        self.carry_model = carry_model
        self.optimize_model = optimize_model
        self.val_model = val_model

        # 保存最近一次的结果，便于复用
        self.evolved_population = None  # 存储最近一次的种群
        self.results_dict = None        # 存储最近一次的结果
        self.val_situation = None       # 存储最近一次的验证集状态
        self.val_results_dict = None    # 存储最近一次的在验证集上的结果
        self.iterative_process = []     # 存储每次迭代的种群得分信息
    
    def get_some_examples_text(self, 
                          examples: Tuple[List[str], List[str]], 
                          examples_number: int) -> Tuple[List[str], List[str]]:
        """
        抽取示例样本（模拟环境作用），用于生成提示词种群
        
        参数:
            examples: 示例训练数据，格式为(输入列表, 输出列表)
            examples_number: 示例样本数量
            
        返回:
            示例文本"输入: {inp}\n输出: {out}"
        """
        assert len(examples[0]) == len(examples[1]), "示例数据长度不一致"
        # 确保抽取的样本数量不超过输入列表和输出列表的长度
        examples_number = min(examples_number, len(examples[0]), len(examples[1]))
        # 将输入和输出组合成一对一对的数据
        paired_examples = list(zip(examples[0], examples[1]))
        # 随机抽取指定数量的样本对
        sampled_paired_examples = random.sample(paired_examples, examples_number)
        # 将抽取的样本对分解为输入和输出列表
        inputs, outputs = zip(*sampled_paired_examples)
        # 准备示例文本
        demo_text = ""
        for i, (inp, out) in enumerate(zip(inputs, outputs)):
            demo_text += f"## 示例{i}\n### 输入: \n{inp}\n### 输出: \n{out}\n\n"
        # 返回示例文本
        return demo_text
        
    def generate_initial_population(self, 
                       examples: Tuple[List[str], List[str]], 
                       examples_number: int,
                       task_description: str = None,
                       population_size: int = 50) -> List[str]:
        """
        抽取示例样本，用于生成初始提示词种群
        
        参数:
            examples: 示例训练数据，格式为(输入列表, 输出列表)
            examples_number: 示例样本数量
            task_description: 任务描述
            population_size: 要生成的提示词数量
            
        返回:
            生成的候选提示词列表: prompts
        """
        # 提示生成模板
        gen_template = PROMPT_generate_TEMPLATE
        # 如果提供了任务描述，则将其添加到提示生成模板中
        if task_description is not None:
            gen_template = gen_template.replace("[TASK]", f"\n\n# 任务描述\n\n{task_description}\n\n")
        else:
            gen_template = gen_template.replace("[TASK]", "")
        
        # 生成候选提示词
        prompts = []
        # 创建进度条
        with tqdm(total=population_size, desc="生成候选提示词") as pbar:
            # 生成指定数量的提示词
            while len(prompts) < population_size:
                # 抽取示例样本
                demo_text = self.get_some_examples_text(examples, examples_number)
                # 准备生成查询
                query = gen_template.replace("[DEMO]", demo_text)

                # 返回的内容
                response = self.optimize_model.communicate(query)
                # 使用正则表达式匹配<prompt>中的内容，包括换行符和特殊字符
                pattern = r'<prompt>(.*?)</prompt>'
                match = re.search(pattern, response, re.DOTALL)  # re.DOTALL使.匹配包括换行符在内的任意字符
                if match:
                    prompt = match.group(1)  # 返回第一个括号中匹配的内容
                else:
                    continue  # 如果没有匹配到，继续
                # 添加到候选列表（不要重复的）
                if prompt and prompt not in prompts:
                    prompts.append(prompt)
                    pbar.update(1) # 更新进度条
        # 返回 候选提示词 列表
        return prompts
    

    def evaluate_prompt(self,
                            prompts: list[str],
                            eval_data_x: list[str],
                            eval_data_y: list[str],
                            eval_config: dict,
                            eval_optimized_config: dict):
        """
        基础评估方法，使用batch_UCB评估提示词
        
        参数:
            prompts: 要评估的提示词列表（不可重复）
            eval_data_x: 评估数据，输入列表
            eval_data_y: 评估数据，输出列表
            eval_config: 评估参数配置
            samples_number: 样本数量
            if_back: 是否放回抽样
            c: 探索参数
            T: 评估总轮数
            
        返回:
            situation: 评估状态
            results: 评估结果，复杂结构的列表
        """
        # 使用batch_UCB进行评估
        try:
            situation, results = batch_UCB(prompts, eval_data_x, eval_data_y, eval_config,
                    eval_optimized_config['samples_number'], eval_optimized_config['if_back'],
                    eval_optimized_config['c'], eval_optimized_config['T'])
        except Exception as e:
            # 这里处理错误
            print(f"An error occurred: {e}")
            # 打印堆栈跟踪信息
            traceback.print_exc()
            print(prompts, eval_data_x, eval_data_y)

        return situation, results
    
    def get_prompt_scores(self, 
                          kind: str, 
                          prompts: list, 
                          prompt_hashs: list,
                          prompt_hash_map: dict,
                          eval_data_x: list, 
                          eval_data_y: list, 
                          eval_config: dict, 
                          eval_optimized_config: dict = None,
                          val_number: int = None,
                          iteration: int = 0):
        """
        获取提示词的得分
        
        输入:
            kind: 评估类型['get_scores', 'best_prompt']
            prompts: 提示词列表
            prompt_hashs: 提示词哈希列表
            prompt_hash_map: 提示词哈希映射
            eval_data_x: 评估数据
            eval_data_y: 评估数据标签
            eval_config: 评估参数配置
            eval_optimized_config: 评估优化参数配置
            val_number: 评估样本数量（评估样本量，因为需要有分数支撑，所以是抽取子集评估）
            iteration: 当前的迭代次数
        
        输出:
            situation: 评估状态
            results_dict: 评估结果（字典，键为提示词哈希，值为列表）
            prompt_scores: 提示词得分（字典，键为哈希，值为平均得分）
        """
        # 评估提示词
        if kind == 'get_scores':
            # 单纯的抽样评估
            eval_optimized_config_get_scores = {                
                'samples_number': val_number, # 样本数量
                'if_back': False,             # 是否放回抽样（默认不放回）
                'c': None,                    # 探索参数
                'T': None,                    # 评估总轮数
            }
            situation, results = self.evaluate_prompt(
                prompts, eval_data_x, eval_data_y,                                          # 基本参数
                eval_config,                                                                # 评估指标参数
                eval_optimized_config_get_scores)
        elif kind == 'best_prompt': 
            # 得到最好的提示词即可
            situation, results = self.evaluate_prompt(
                prompts, eval_data_x, eval_data_y,                                          # 基本参数
                eval_config,                                                                # 评估指标参数
                eval_optimized_config)
        else:
            raise ValueError(f"未知的评估类型: {kind}")
        # 计算每个提示词的得分
        prompt_scores = {}
        # 评估结果字典（键为提示词哈希，值为列表）
        results_dict = {}
        # results [[prompt, combined_data, average_score]，，，，，]
        # combined_data: (x, y_real, y_pred_original, y_pred, 得分, 评估解释)
        for item in results:
            prompt, combined_data, avg_score = item
            prompt_hash = hash(prompt)
            prompt_scores[prompt_hash] = avg_score
            results_dict[prompt_hash] = [prompt, combined_data, avg_score, iteration]
        # 返回结果
        return situation, results_dict, prompt_scores
    

#################################################### 反馈进化相关 ###########################################
#################################################### 反馈进化相关 ###########################################
#################################################### 反馈进化相关 ###########################################
    def _select_parents(self, population, prompt_scores, parents_number, select_kind):
        """
        选择亲代
        
        参数:
            population: 当前提示词种群（哈希形式的列表）
            prompt_scores: 提示词得分字典，键为提示词，值为得分
            parents_number: 要选择的亲代数量
            select_kind: 选择方法，可以是 'wheel'（轮盘赌）, 'nr_wheel'（非重复轮盘赌）, 'random'（随机）
            
        返回:
            选择的亲代列表（哈希形式的列表）
        """
        if select_kind == 'random':
            # 随机选择亲代
            return random.sample(population, min(parents_number, len(population)))
        elif select_kind == 'wheel':
            # 轮盘赌选择亲代
            weights = [prompt_scores.get(prompt, 0) for prompt in population]
            total_weight = sum(weights)
            # 如果所有权重都不是正数，则随机选择
            if total_weight <= 0:
                return random.sample(population, min(parents_number, len(population)))
            # 依次放回抽取 parents_number 个亲代parents
            parents = []
            for _ in range(min(parents_number, len(population))):
                norm_weights = [w / total_weight for w in weights]
                idx = random.choices(range(len(population)), weights=norm_weights, k=1)[0]
                parents.append(population[idx])
            return parents
        elif select_kind == 'nr_wheel':
            # 非重复轮盘赌选择亲代
            weights = [prompt_scores.get(prompt, 0) for prompt in population]
            total_weight = sum(weights)
            # 如果所有权重都不是正数，则随机选择
            if total_weight <= 0:
                return random.sample(population, min(parents_number, len(population)))
            # 依次不放回抽取 parents_number 个亲代parents
            parents = []
            available_indices = list(range(len(population)))
            available_weights = weights.copy()
            # 循环选择 parends_number 个亲代
            for _ in range(min(parents_number, len(population))):
                total_available_weight = sum(available_weights)
                norm_weights = [w / total_available_weight for w in available_weights]
                idx = random.choices(available_indices, weights=norm_weights, k=1)[0]
                # 添加到亲代列表并从可用列表中移除
                parents.append(population[idx])
                idx_to_remove = available_indices.index(idx)
                available_indices.pop(idx_to_remove)
                available_weights.pop(idx_to_remove)
            return parents
        else:
            raise ValueError(f"未知的选择方法: {select_kind}")
    
    def _sample_examples(self, 
                         train_number:  int, 
                         sampling_method: str, 
                         parents_hash: list[str], 
                         results_dict: Dict[str, list]):
        """
        对于每个亲代，抽取样本
        
        参数:
            train_number: 训练样本数量（要抽取的样本数量）
            sampling_method: 采样方法
            parents_hash: 亲代提示词哈希
            results_dict: 评估结果（字典，键为提示词哈希，值为列表）
        返回:
            samples: 每个提示词抽取的样本字典
        """
        samples = {}
        # 遍历每个亲代，进行采样
        for parent_hash in parents_hash:
            # 提取各个样本的情况 [prompt, combined_data, average_score]
            # combined_data: (x, y_real, y_pred_original, y_pred, 得分, 评估解释)
            combined_data = results_dict[parent_hash][1]
            # 从大到小排序
            sorted_samples = sorted(combined_data, key=lambda x: x[4] if x[4] is not None else float('-inf'), reverse=True)

            if sampling_method == 'random':
                # 随机采样
                parent_samples = random.sample(sorted_samples, min(train_number, len(combined_data)))
                samples[parent_hash] = parent_samples

            elif sampling_method == 'balanced':
                # 平衡采样（选择一半表现好的样本和一半表现差的样本）
                sorted_samples = sorted(combined_data, key=lambda x: x[4])
                half_size = train_number // 2
                parent_samples = [ii for ii in sorted_samples[:half_size]]  # 表现差的样本
                parent_samples += [ii for ii in sorted_samples[-half_size:]]  # 表现好的样本
                samples[parent_hash] = parent_samples

            elif sampling_method == 'advantage':
                # 优势采样（选择表现好的样本）
                sorted_samples = sorted(combined_data, key=lambda x: x[4], reverse=True)
                parent_samples = [ii for ii in sorted_samples[:train_number]]
                samples[parent_hash] = parent_samples

            elif sampling_method == 'disadvantage':
                # 劣势采样（选择表现差的样本）
                sorted_samples = sorted(combined_data, key=lambda x: x[4])
                parent_samples = [ii for ii in sorted_samples[:train_number]]
                samples[parent_hash] = parent_samples

            else:
                raise ValueError(f"未知的采样方法: {sampling_method}")
        # 返回选择的样本
        return samples
    

    def _prepare_demo_text(self, parents_hash, prompt_hash_map, sampled_examples):
        '''
        准备示例文本
        
        参数:
            parents_hash: 亲代哈希列表
            prompt_hash_map: 提示词哈希映射（字典，键为提示词哈希，值为提示词）
            sampled_examples: 每个亲代抽取的样本字典
        返回:
            demo_text: 示例文本
        '''
        demo_text = ""
        for i, now_hash in enumerate(parents_hash):
            demo_text += f'''
## 提示词{i}
{prompt_hash_map[now_hash]}

## 示例上的表现{i}
'''
            for j, (x, y_real, y_pred_original, y_pred, score, reason) in enumerate(sampled_examples[now_hash]):
                demo_text += f'''
### 示例上的表现{j}

#### 输入: 
{x}

#### 预期输出: 
{y_real}

#### 模型实际的原始输出: 
{y_pred_original}

#### 模型提取后的应用输出: 
{y_pred}

#### 得分: 
{score}

'''
                if reason is None:
                    pass
                else:
                    demo_text += f'''
#### 得分解释:
{reason}

'''
            demo_text += "\n\n"
        # 返回
        return demo_text

    def _generate_feedback_base(self, 
                           parents_hash: List[str], 
                           prompt_hash_map: Dict[str, str], 
                           sampled_examples: Dict[str, List[Tuple[str, str, str, str, float, str]]],
                           task_description: str = None,):
        """
        生成反馈
        
        参数:
            parents_hash: 亲代提示词哈希
            prompt_hash_map:  提示词哈希映射
            sampled_examples: 抽取的样本字典，键为哈希,值为列表，(x, y_real, y_pred_original, y_pred, 得分, 评估解释)
            task_description: 任务描述
        返回:
            feedback: 为提示词生成的反馈
        """
        # 准备示例文本（与原代码相同）
        demo_text = self._prepare_demo_text(parents_hash, prompt_hash_map, sampled_examples)
        
        
        
        # 提示生成模板
        gen_template = PROMPT_feedback_TEMPLATE
        # 如果提供了任务描述，则将其添加到提示生成模板中
        if task_description is not None:
            gen_template = gen_template.replace("[TASK]", f"\n\n# 任务描述\n\n{task_description}\n\n")
        else:
            gen_template = gen_template.replace("[TASK]", "")

        # 准备生成查询
        query = gen_template.replace("[PROMPT and DEMO]", demo_text)
        print("\n\生成反馈\n\n")
        print(query)
        # 获取反馈
        response = self.optimize_model.communicate(query)
        
        # 使用正则表达式匹配<feedback>中的内容
        pattern = r'<feedback>(.*?)</feedback>'
        match = re.search(pattern, response, re.DOTALL)
        
        if match:
            feedback = match.group(1).strip()
            return feedback
        else:
            # 如果没有找到匹配的标签，则返回整个响应
            return response
        


    def _generate_feedback_strong(self, 
                       parents_hash: List[str], 
                       prompt_hash_map: Dict[str, str], 
                       sampled_examples: Dict[str, List[Tuple[str, str, str, str, float, str]]],
                       task_description: str = None):
        """
        增强的多层次反馈生成机制（先从 错误模式识别, 根因分析, 改进建议生成 三个方面，然后综合为反馈）

        参数:
            parents_hash: 亲代提示词哈希
            prompt_hash_map:  提示词哈希映射
            sampled_examples: 抽取的样本字典，键为哈希,值为列表，(x, y_real, y_pred_original, y_pred, 得分, 评估解释)
            task_description: 任务描述
        返回:
            feedback: 为提示词生成的反馈
        """
        # 准备示例文本
        demo_text = self._prepare_demo_text(parents_hash, prompt_hash_map, sampled_examples)
        
        # 1. 错误模式识别
        error_patterns_prompt = PROMPT_error_patterns_TEMPLATE.replace("[PROMPT and DEMO]", demo_text)
        if task_description is not None:
            error_patterns_prompt = error_patterns_prompt.replace("[TASK]", f"\n\n# 任务描述\n\n{task_description}\n\n")
        else:
            error_patterns_prompt = error_patterns_prompt.replace("[TASK]", "")
        
        error_patterns_response = self.optimize_model.communicate(error_patterns_prompt)
        error_patterns = self._extract_tag_content(error_patterns_response, "error_patterns")
        
        # 2. 根因分析
        root_cause_prompt = PROMPT_root_cause_TEMPLATE.replace("[ERROR_PATTERNS]", error_patterns)
        if task_description is not None:
            root_cause_prompt = root_cause_prompt.replace("[TASK]", f"\n\n# 任务描述\n\n{task_description}\n\n")
        else:
            root_cause_prompt = root_cause_prompt.replace("[TASK]", "")
        
        root_cause_response = self.optimize_model.communicate(root_cause_prompt)
        root_causes = self._extract_tag_content(root_cause_response, "root_causes")
        
        # 3. 改进建议生成
        improvement_prompt = PROMPT_improvement_TEMPLATE.replace("[ROOT_CAUSES]", root_causes)
        if task_description is not None:
            improvement_prompt = improvement_prompt.replace("[TASK]", f"\n\n# 任务描述\n\n{task_description}\n\n")
        else:
            improvement_prompt = improvement_prompt.replace("[TASK]", "")
        
        improvement_response = self.optimize_model.communicate(improvement_prompt)
        improvements = self._extract_tag_content(improvement_response, "improvements")
        
        # 综合反馈
        feedback_prompt = PROMPT_feedback_synthesis_TEMPLATE.replace(
            "[ERROR_PATTERNS]", error_patterns).replace(
            "[ROOT_CAUSES]", root_causes).replace(
            "[IMPROVEMENTS]", improvements)
        
        feedback_response = self.optimize_model.communicate(feedback_prompt)
        feedback = self._extract_tag_content(feedback_response, "feedback")
        
        if feedback:
            return feedback
        else:
            return feedback_response
        
    def _extract_tag_content(self, text, tag_name):
        """从文本中提取指定标签的内容"""
        pattern = f'<{tag_name}>(.*?)</{tag_name}>'
        match = re.search(pattern, text, re.DOTALL)
        if match:
            return match.group(1).strip()
        return ""


    def _extract_knowledge(self, feedback):
        """从反馈中提取领域知识"""
        knowledge_prompt = PROMPT_knowledge_extraction_TEMPLATE.replace("[FEEDBACK]", feedback)
        knowledge_response = self.optimize_model.communicate(knowledge_prompt)
        # 解析知识为结构化格式
        pattern = r'<{knowledge}>(.*?)</{knowledge}>'
        matches = re.findall(pattern, knowledge_response, re.DOTALL)
        # 使用列表推导式去除每个匹配项的前后空白字符
        knowledge_items = [match.strip() for match in matches]
        
        return knowledge_items

    def _generate_children(self, 
                           parents_hash: List[str], 
                           prompt_hash_map: Dict[str, str], 
                           examples: Tuple[List[str], List[str]], 
                           examples_number: int,
                           feedback: str, 
                           task_description: str):
        """
        生成子代
        
        参数:
            parents_hash: 提示词哈希
            prompt_hash_map:  提示词哈希映射
            examples: 示例训练数据，格式为(输入列表, 输出列表)
            examples_number: 示例样本数量
            feedback: 为提示词生成的反馈
            task_description: 任务描述
            
        返回:
            child_prompts: 生成的子代提示词列表
        """
        # 生成的子代
        child_prompts = []
        # 改进方向（增加/删减/修改）
        improved_direction = ["增加", "删减", "修改"]
        # 准备亲代提示词文本
        prompt_text = ""
        for i, now_hash in enumerate(parents_hash):
            prompt_text += f'''
## 提示词{i}
{prompt_hash_map[now_hash]}

'''

        # 沿方向改进
        for now_direct in improved_direction:
            # 提示生成模板
            gen_template = PROMPT_child_TEMPLATE
            # 如果提供了任务描述，则将其添加到提示生成模板中
            if task_description is not None:
                gen_template = gen_template.replace("[TASK]", f"\n\n# 任务描述\n\n{task_description}\n\n")
            else:
                gen_template = gen_template.replace("[TASK]", "")
            # 映射进化算子
            now_direct_info = PROMPT_direct_TEMPLATE[now_direct]
            # 抽取示例样本
            demo_text = self.get_some_examples_text(examples, examples_number)
                
            # 准备生成查询 
            query = gen_template.replace("[PROMPT]", prompt_text).replace("[DEMO]", demo_text)
            query = query.replace("[FEEDBACK]", feedback).replace("[DIRECTION]", now_direct_info)
            print("\n\n生成子代\n\n")
            print(query)
            # 生成子代
            response = self.optimize_model.communicate(query)
            
            # 使用正则表达式匹配<prompt>中的内容
            pattern = r'<prompt>(.*?)</prompt>'
            match = re.search(pattern, response, re.DOTALL)
            if match:
                child_prompt = match.group(1).strip()
                if child_prompt and child_prompt not in child_prompts:
                    child_prompts.append(child_prompt)
        
        return child_prompts
    

    def evolve_prompts(self,
                        examples: Tuple[List[str], List[str]], 
                        eval_config: dict,
                        eval_optimized_config: dict,

                        population: List[str],
                        prompt_hashs: List[str],
                        prompt_hash_map: Dict[str, str],
                        prompt_scores: Dict[str, Any],
                        results_dict: Dict[str, list],

                        examples_number: int = 5,
                        task_description: str = None,
                        num_iterations: int = 5,
                        population_size: int = 50,
                        parents_number: int = 10,
                        select_kind: str = "wheel",
                        val_number: int = 10,
                        train_number: int = 10,
                        sampling_method: str = "random",
                        feedback_kind: str = "base",
                        knowledge_kind: bool = False):
        """
       反馈进化
        
        基本参数:
            examples: 训练数据，格式为(输入列表, 输出列表)
            eval_config: 评估参数配置
            eval_optimized_config: 评估优化参数配置

            population: 初始提示词种群
            prompt_hashs: 提示词哈希列表
            prompt_hash_map:  提示词哈希映射
            prompt_scores: 提示词得分字典，键为提示词哈希，值为得分
            results_dict: 评估结果（字典，键为提示词哈希，值为列表）

            examples_number: 示例样本数量（用于生成提示词）
            task_description: 任务描述
            num_iterations: 迭代次数
            population_size: 种群数量
            parents_number: 亲代数量
            select_kind: 亲代选择方法['wheel', 'nr_wheel', 'random']（轮盘赌/非重复轮盘赌/随机）
            val_number: 评估样本数量（用于评估子代、更新种群）
            train_number: 训练样本数量（用于抽取样本、生成反馈）
            sampling_method: 采样方法['random', 'balanced', 'advantage', 'disadvantage'](随机采样/平衡采样/优势采样/劣势采样)
            feedback_kind: 反馈生成模式['base', 'strong'](基础/增强)
            knowledge_kind: 是否使用知识库积累[True, False]

        评估指标参数：
            eval_config: 一个字典，用于评估指标的具体参数
        batch-UCB 评估优化参数（可选）：默认穷举
            samples_number: 样本数量
            if_back: 是否放回抽样（默认不放回）
            c: 探索参数（默认为0）
            T: 评估总轮数（默认为待评估对象数量）
        返回:
            prompt_hash_map:  提示词哈希映射
            prompt_scores: 提示词得分（字典，键为提示词哈希，值为得分）
            results_dict: 评估结果（字典，键为提示词哈希，值为列表）
            current_population_prompts: 当前种群提示词
        """
        # 当前种群（不断更新）
        current_population_hash = prompt_hashs.copy()
        current_population_prompts = population.copy()
        # 下面三个算是全局变量记录
        # prompt_hash_map:  提示词哈希映射
        # prompt_scores: 提示词得分字典，键为提示词哈希，值为得分
        # results_dict: 评估结果（字典，键为提示词哈希，值为列表）
        # 初始化知识库
        knowledge_base = set()

        # 创建进度条
        with tqdm(total=num_iterations, desc="FE进化迭代（轮数）") as pbar:
            for iteration in range(num_iterations):
                ######### 3.1 选择亲代
                parents_hash = self._select_parents(current_population_hash, prompt_scores, parents_number, select_kind)
                
                ######### 3.2 抽取样本（从亲代的已评估样本中采样）
                sampled_examples = self._sample_examples(train_number, sampling_method, parents_hash, results_dict)
                
                ######### 3.3 生成反馈
                if feedback_kind == 'strong':
                    feedback = self._generate_feedback_strong(parents_hash, prompt_hash_map, sampled_examples, task_description)
                else:
                    feedback = self._generate_feedback_base(parents_hash, prompt_hash_map, sampled_examples, task_description)
                
                # 使用积累的知识库增强
                if knowledge_kind == True:
                    # 提取和累积知识
                    new_knowledge = self._extract_knowledge(feedback)
                    knowledge_base.update(new_knowledge)
                    # 将累积的知识融入子代生成
                    enhanced_feedback = feedback
                    if knowledge_base:
                        knowledge_text = "\n".join(list(knowledge_base))
                        enhanced_feedback += f"\n\n# 累积的领域知识\n\n{knowledge_text}"
                    feedback = enhanced_feedback

                ######### 3.4 生成子代（增加/删减/修改）
                # 生成的新提示词
                new_population = []
                new_population_hash = []
                child_prompts = self._generate_children(parents_hash, prompt_hash_map, 
                                                        examples,examples_number,
                                                        feedback, task_description)
                # 添加到新种群（确保不重复）
                for child in child_prompts:
                    child_hash = hash(child)
                    if child_hash not in current_population_hash and child not in new_population:
                        new_population.append(child)
                        new_population_hash.append(child_hash)
                        prompt_hash_map[child_hash] = child
                
                ######### 3.5 评估子代
                if new_population:
                    # 评估新子代（用训练集）
                    eval_data_x, eval_data_y = examples[0], examples[1]
                    # 获取提示词及其得分(new_results_dict、prompt_scores的键为提示词哈希)
                    situation, new_results_dict, new_prompt_scores = self.get_prompt_scores(
                        'get_scores', new_population, prompt_hashs, prompt_hash_map, eval_data_x, eval_data_y, \
                            eval_config, eval_optimized_config, val_number, iteration+1)
                    # 更新结果和得分
                    results_dict.update(new_results_dict)
                    prompt_scores.update(new_prompt_scores)
                
                ######### 3.6 更新种群（择优）
                # 合并亲代和子代
                combined_population_hash = current_population_hash + new_population_hash
                # 按照得分进行从大到小
                sorted_population_hash = sorted(combined_population_hash, key=lambda p: prompt_scores.get(p, 0), reverse=True)
                # 选择前 population_size 个最佳提示词哈希
                current_population_hash = sorted_population_hash[:population_size]
                current_population_prompts = [prompt_hash_map[p] for p in current_population_hash]

                ######### 记录当前迭代的种群得分信息
                current_scores = [prompt_scores.get(p, 0) for p in current_population_hash]
                iteration_info = {
                    "iteration": iteration + 1,
                    "population_size": len(current_population_hash),
                    "average_score": sum(current_scores) / len(current_scores) if current_scores else 0,
                    "max_score": max(current_scores) if current_scores else 0,
                    "min_score": min(current_scores) if current_scores else 0,
                    "scores": current_scores.copy()
                }
                self.iterative_process.append(iteration_info)
                # 进度条更新
                pbar.update(1)
        
        # 返回进化后的种群
        return prompt_hash_map, prompt_scores, results_dict, current_population_prompts
    
#################################################### 保存 ###########################################
#################################################### 保存 ###########################################
#################################################### 保存 ###########################################
    def _save_results(  self,
                        evolved_population, 
                        results_dict, 
                        prompt_hash_map,
                        val_situation, 
                        val_results_dict,
                        examples,
                        val_examples,
                        eval_config,
                        eval_optimized_config,

                        initial_population,
                        examples_number,
                        task_description,
                        num_iterations,
                        population_size,
                        parents_number,
                        select_kind,
                        val_number,
                        train_number,
                        sampling_method,
                        feedback_kind, 
                        knowledge_kind,

                        highest_prompt,
                        highest_score):
        '''保存结果（输入、参数配置、结果）'''
        # 获取当前文件的所在文件夹路径
        current_file_directory = os.path.dirname(os.path.abspath(__file__))
        # 结果文件夹
        folder_name = os.path.join(current_file_directory, "simplified_FE_results")
        if not os.path.exists(folder_name):
            os.makedirs(folder_name)
        # 本次运行的结果的文件夹（以结束时间命名）
        current_time = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
        folder_name = os.path.join(folder_name, f"time-{current_time}")
        if not os.path.exists(folder_name):
            os.makedirs(folder_name)

        # 训练集
        # 评估结果 results 转换为 JSON 兼容的格式
        json_results = []
        for key_hash, item in results_dict.items():
            prompt, samples, avg_score, iteration = item
            json_samples = []
            for sample in samples:
                x, y_real, y_pred_original, y_pred, score, reason = sample
                json_sample = {
                    "x": x,
                    "y_real": y_real,
                    "y_pred_original": y_pred_original,
                    "y_pred": y_pred,
                    "score": score,
                    "reason": reason
                }
                json_samples.append(json_sample)
            json_item = {
                "prompt": prompt,
                "samples": json_samples,
                "average_score": avg_score,
                "iteration": iteration
            }
            json_results.append(json_item)
        # 排序（按照 average_score 从大到小）
        sorted_json_results = sorted(json_results, key=lambda item: item["average_score"], reverse=True)
        # 将转换后的数据保存到 JSON 文件
        json_file_path = os.path.join(folder_name, 'results.json')
        with open(json_file_path, 'w', encoding='utf-8') as f:
            json.dump(sorted_json_results, f, ensure_ascii=False, indent=4)

        # 验证集
        # 评估结果 results 转换为 JSON 兼容的格式
        json_results = []
        for key_hash, item in val_results_dict.items():
            prompt, samples, avg_score, iteration = item
            json_samples = []
            for sample in samples:
                x, y_real, y_pred_original, y_pred, score, reason = sample
                json_sample = {
                    "x": x,
                    "y_real": y_real,
                    "y_pred_original": y_pred_original,
                    "y_pred": y_pred,
                    "score": score,
                    "reason": reason
                }
                json_samples.append(json_sample)
            json_item = {
                "prompt": prompt,
                "samples": json_samples,
                "average_score": avg_score,
                "iteration": iteration
            }
            json_results.append(json_item)
        # 排序（按照 average_score 从大到小）
        sorted_json_results = sorted(json_results, key=lambda item: item["average_score"], reverse=True)
        # 将转换后的数据保存到 JSON 文件
        json_file_path = os.path.join(folder_name, 'results_val.json')
        with open(json_file_path, 'w', encoding='utf-8') as f:
            json.dump(sorted_json_results, f, ensure_ascii=False, indent=4)

        # 创建 info.txt 文件
        info_file_path = os.path.join(folder_name, 'info.txt')
        with open(info_file_path, 'w', encoding='utf-8') as f:
            f.write(f"时间: {current_time}\n")
            f.write(f"训练集大小: {len(examples[0])}\n")
            f.write(f"验证集大小: {len(val_examples[0])}\n")
            f.write(f"初始提示词数量: {len(initial_population)}\n")
            f.write(f"示例样本数量（用于生成初始提示词）: {examples_number}\n")
            f.write(f"任务描述: {task_description}\n")
            f.write(f"迭代次数: {num_iterations}\n")
            f.write(f"种群数量: {population_size}\n")
            f.write(f"亲代数量: {parents_number}\n")
            f.write(f"选择方法: {select_kind}\n")
            f.write(f"评估样本数量（用于评估子代、更新种群）: {val_number}\n")
            f.write(f"训练样本数量（用于抽取样本、生成反馈）: {train_number}\n")
            f.write(f"采样方法: {sampling_method}\n")
            f.write(f"反馈方法: {feedback_kind}\n")
            f.write(f"知识库累积: {knowledge_kind}\n")
            f.write(f"\n")
            f.write(f"最后一代种群: {evolved_population}\n")
            f.write(f"验证状态: {val_situation}\n")
            f.write(f"得分最高的提示词: {highest_prompt}\n")
            f.write(f"最高分: {highest_score}\n")

        # 实例无法被 json.dump 序列化，需要替换掉
        # 强制将所有不可序列化的对象转换为字符串
        def force_stringify(obj):
            if isinstance(obj, (dict, list, str, int, float, bool, type(None))):
                return obj
            else:
                return str(obj)
        # eval_config: 评估参数配置
        json_file_path = os.path.join(folder_name, 'eval_config.json')
        with open(json_file_path, 'w', encoding='utf-8') as f:
            json.dump(eval_config, f, default=force_stringify, ensure_ascii=False, indent=4)
        # eval_optimized_config: 评估优化配置
        json_file_path = os.path.join(folder_name, 'eval_optimized_config.json')
        with open(json_file_path, 'w', encoding='utf-8') as f:
            json.dump(eval_optimized_config, f, default=force_stringify, ensure_ascii=False, indent=4)

        # 保存迭代过程信息
        json_file_path = os.path.join(folder_name, 'iterative_process.json')
        with open(json_file_path, 'w', encoding='utf-8') as f:
            json.dump(self.iterative_process, f, ensure_ascii=False, indent=4)
        return 0

    

    def find_best_prompt(self,
                        examples: Tuple[List[str], List[str]], 
                        val_examples: Tuple[List[str], List[str]], 
                        eval_config: dict,
                        eval_optimized_config: dict,

                        initial_population: List[str] = [],
                        examples_number: int = 5,
                        task_description: str = None,
                        num_iterations: int = 5,
                        population_size: int = 50,
                        parents_number: int = 10,
                        select_kind: str = "wheel",
                        val_number: int = 10,
                        train_number: int = 10,
                        sampling_method: str = "random",
                        feedback_kind: str = "base",
                        knowledge_kind: bool = False):
        """
        核心功能：自动找到最佳提示词（一体化的函数）
        
        参数:
            基本参数:
            examples: 训练数据，格式为(输入列表, 输出列表)
            val_examples: 验证数据，格式为(输入列表, 输出列表)
            eval_config: 评估参数配置
            eval_optimized_config: 评估优化参数配置

            initial_population: 初始提示词列表 
            examples_number: 示例样本数量（用于生成初始提示词）
            task_description: 任务描述
            num_iterations: 迭代次数
            population_size: 种群数量
            parents_number: 亲代数量
            select_kind: 亲代选择方法（轮盘赌/非重复轮盘赌/随机）['wheel', 'nr_wheel', 'random']
            val_number: 评估样本数量（用于评估子代、更新种群）
            train_number: 训练样本数量（用于抽取样本、生成反馈）
            sampling_method: 采样方法(随机采样/平衡采样/优势采样/劣势采样)['random', 'balanced', 'advantage', 'disadvantage']
            feedback_kind: 反馈生成模式['base', 'strong'](基础/增强)
            knowledge_kind: 知识库累积[True, False]
            
        返回:
            situation：评估状态
            best_prompt：最佳提示词

        """
        ######### 1. 生成初始提示词种群(如果传入的数量不够，则自动生成)（使用训练集）
        print("1. 生成初始提示词种群...")
        if initial_population is None:
            # 没有就从头生成
            initial_population = self.generate_initial_population(examples, examples_number, task_description, population_size)
        elif len(initial_population) < population_size:
            # 差多少补多少
            initial_population_new = self.generate_initial_population(examples, examples_number, task_description, population_size - len(initial_population))
            # 合并
            initial_population += initial_population_new
        # 创建哈希映射，用于简化提示词表示
        prompt_hashs = [hash(prompt) for prompt in initial_population]
        prompt_hash_map = {hash(prompt): prompt for prompt in initial_population}

        ######### 2. 评估初始提示词种群（使用验证集）
        print("2. 评估初始提示词种群...")
        eval_data_x, eval_data_y = examples[0], examples[1]
        # 获取提示词及其得分(results中为原始提示词，prompt_scores的键为提示词哈希)
        situation, results_dict, prompt_scores = self.get_prompt_scores(
            'get_scores', initial_population, 
            prompt_hashs, prompt_hash_map, 
            eval_data_x, eval_data_y, 
            eval_config, eval_optimized_config, val_number, 0)
            
        ######### 记录初始种群得分信息
        self.iterative_process = []  # 重置追踪信息
        initial_scores = [prompt_scores.get(h, 0) for h in prompt_hashs]
        initial_info = {
            "iteration": 0,
            "population_size": len(prompt_hashs),
            "average_score": sum(initial_scores) / len(initial_scores) if initial_scores else 0,
            "max_score": max(initial_scores) if initial_scores else 0,
            "min_score": min(initial_scores) if initial_scores else 0,
            "scores": initial_scores.copy()
        }
        self.iterative_process.append(initial_info)

        ######### 3. 使用FE算法进行反馈进化
        print("3. 使用FE算法进行反馈进化...")
        prompt_hash_map, prompt_scores, results_dict, evolved_population = self.evolve_prompts(
            examples, eval_config, eval_optimized_config,
            initial_population, prompt_hashs, prompt_hash_map, prompt_scores, results_dict,
            examples_number, task_description, num_iterations, population_size, parents_number,
            select_kind, val_number, train_number, sampling_method, feedback_kind, knowledge_kind
        )

        ######### 4. 评估找到最佳提示词（验证集上）
        print("4. 评估找到最佳提示词...")
        eval_data_x, eval_data_y = val_examples[0], val_examples[1]
        val_situation, val_results_dict, val_prompt_scores = self.get_prompt_scores(
            'best_prompt', evolved_population, 
            prompt_hashs, prompt_hash_map, 
            eval_data_x, eval_data_y, 
            eval_config, eval_optimized_config, val_number, num_iterations+1
        )
        # 找到最佳提示词
        highest_score = -float('inf')
        highest_prompt_hash = None
        for now_hash, avg_score in val_prompt_scores.items():
            if avg_score > highest_score:
                highest_score = avg_score
                highest_prompt_hash = now_hash
        # 由哈希映射回提示词
        highest_prompt = prompt_hash_map[highest_prompt_hash]
        print(f"最佳提示词: {highest_prompt}")
        print(f"最佳得分: {highest_score}")
        # 保存最近一次的结果，便于复用
        self.evolved_population = evolved_population
        self.results_dict = results_dict
        self.val_situation = val_situation
        self.val_results_dict = val_results_dict

        ######### 4. 保存过程与结果
        print("6. 保存结果...")
        self._save_results(evolved_population, results_dict, prompt_hash_map,
                           val_situation, val_results_dict,
                           examples, val_examples, eval_config, eval_optimized_config,
                        initial_population, examples_number, task_description, num_iterations,
                        population_size, parents_number, select_kind, val_number, 
                        train_number, sampling_method, feedback_kind, knowledge_kind,
                        highest_prompt, highest_score)

        return situation, highest_prompt, val_prompt_scores












# 与 执行模版 配套的 执行模式匹配器
class CarryMatcher(BaseMatcher):
    def matcher(self, input_string):
        """
        重写基类的matchor方法，实现具体的匹配逻辑

        输入：
            input_string: 大模型的原始生成内容，即要匹配的字符串
        输出：
            [input_string]，匹配内容的列表
        """
        
        # 不进行任何匹配

        return [input_string]
    


# 演示用法
def main():
    """主函数：演示APE的使用"""
    
    
    # 创建模型和APE实例
    model = Da_Mo_Xing_2("Qwen2.5-7B-Instruct")
    FE = SimpleFE(model, model, model)
    
    # 准备示例数据
    words = ["聪明", "直接", "正式地", "不受欢迎的", "减法"]
    antonyms = ["愚蠢", "间接", "非正式地", "受欢迎的", "加法"]

    # 执行模板（让 执行模型 根据 提示词、输入生成输出）
    carry_template = "指令: [PROMPT]\n输入: [INPUT]\n输出: [OUTPUT]"
    # 与 执行模版 配套的 执行模式匹配器
    carry_matcher = CarryMatcher()

    # 评估配置
    eval_config = {
            'evaluation_index': 'acc',             # 评估指标
            'carry_template': carry_template,      # 执行模版
            'carry_matcher': carry_matcher,        # 执行模式匹配器，与 执行模版 配套
            'carry_model': model        # 执行模型
        }
    # 评估优化配置
    eval_optimized_config = {
        'samples_number': None, # 样本数量
        'if_back': False, # 是否放回抽样（默认不放回）
        'c': None, # 探索参数
        'T': None, # 评估总轮数
    }
    
    # 找到最佳提示词
    situation, best_prompt, val_avg_score = FE.find_best_prompt(
        (words, antonyms),      # 训练集 (输入列表, 输出列表)
        (words, antonyms),      # 验证集 (输入列表, 输出列表)
        eval_config,            # 评估配置
        eval_optimized_config,  # 评估优化配置
        None,                   # 初始提示词列表（None表示自动生成）
        3,                      # 示例样本数量
        None,                   # 任务描述
        2,                      # 迭代次数
        5,                      # 种群数量
        3,                      # 亲代数量
        "wheel",                # 选择方法
        5,                      # 评估样本数量
        3,                      # 训练样本数量
        "random",               # 采样方法
        "base",                 # 反馈生成方法
        False                   # 是否使用知识库积累
    )
    
    # 演示函数
    def demo_fn(prompt, inputs):
        """测试提示词在新输入上的效果"""
        if not isinstance(inputs, list):
            inputs = [inputs]
        results = []
        carry_model = Da_Mo_Xing_2("Qwen2.5-7B-Instruct")
        for input_text in inputs:
            query = carry_template.replace('[PROMPT]', prompt).replace('[INPUT]', input_text).replace('[OUTPUT]', '')
            # 获取大模型的原始内容
            now_y_pred_original = carry_model.communicate(query)
            # 进行格式解析
            now_y_pred = carry_matcher.matcher(now_y_pred_original)[0]
            results.append(now_y_pred)
        return results
    # 测试最佳提示词
    test_words = ["高", "热", "硬", "新"]
    print("\n测试最佳提示词:")
    results = demo_fn(best_prompt, test_words)
    for word, result in zip(test_words, results):
        print(f"单词: {word}, 反义词: {result}")


if __name__ == "__main__":
    main()